
import re

from my_py_toolkit.file.excel.excel_toolkit import write_excel



def handle_metric(data_path, save_path):
    titles = []
    values = []
    reg_replace = '.*- __main__ - '
    reg_content = 'Dev\\n-{80}\\n(?P<content>[\\s\\S]*?)\\n-{80}\\n\*\*'
    data = ''
    with open(p, 'r', encoding='utf-8') as f:
        data = f.read()
    
    data = re.sub(reg_replace, '', data)
    for match in re.finditer(reg_content, data):
        cur_titles = []
        cur_values = []
        for line in match.group('content').split('\n'):
            k,v = [s.strip() for s in line.split('=')]
            cur_titles.append(k)
            cur_values.append(v)
        titles = cur_titles if not titles else titles
        values.append(cur_values)
    titles, values
    write_excel([titles] + values, save_path)    

def handle_loss(data_path, save_path, epochs=20):
    info = {}
    reg = 'global_steps (?P<global_steps>\d+) - lr: (?P<lr>\d+\.\d+) - loss: (?P<loss>(-)?\d+\.\d+)'
    with open(data_path, 'r', encoding='utf-8') as f:
        for line in f.read().split('\n'):
            match = re.search(reg, line)
            if match:
                for k,v in match.groupdict().items():
                    if k not in info:
                        info[k] = []
                    info[k].append(float(v))
    res = [['min', 'max', 'avg']]
    avg_lens = len(info['loss']) // epochs
    for i in range(epochs):
        loss = info['loss'][i*avg_lens:(i+1)*avg_lens]
        res.append([min(loss), max(loss), sum(loss)/len(loss)])
    write_excel(res, save_path)

if __name__ == '__main__':    
    p = 'outputs/bert-resources_bert_model_bert/run.log'
    handle_metric(p, './tools/handle_metrics.xlsx')
    handle_loss(p, './tools/handle_loss.xlsx')